{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5cf0daf9",
   "metadata": {},
   "source": [
    "# Deep Research with a Clarifying Agent\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5dad4e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from agents import Agent, WebSearchTool, trace, Runner, gen_trace_id, function_tool\n",
    "from agents.model_settings import ModelSettings\n",
    "from pydantic import BaseModel, Field\n",
    "from dotenv import load_dotenv\n",
    "import asyncio\n",
    "import requests\n",
    "import os\n",
    "from typing import Dict\n",
    "from IPython.display import display, Markdown\n",
    "import gradio as gr\n",
    "from typing import List, Tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5998ae2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6870327",
   "metadata": {},
   "source": [
    "### Create an Agent to ask clarifying questions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddd16501",
   "metadata": {},
   "outputs": [],
   "source": [
    "CLARIFYING_INSTRUCTIONS = \"\"\"You are a research analyst. Given a research query, generate 3-5 clarifying \n",
    "questions that would help narrow down and improve the research. Focus on:\n",
    "- Scope and timeframe\n",
    "- Specific aspects of interest\n",
    "- Intended use case or audience\n",
    "- Level of detail needed\n",
    "\n",
    "Make questions concise and easy to answer.\"\"\"\n",
    "\n",
    "class ClarifyingQuestion(BaseModel):\n",
    "    question: str = Field(description=\"The clarifying question to ask\")\n",
    "    why_important: str = Field(description=\"Brief explanation of why this question matters\")\n",
    "\n",
    "class ClarifyingQuestions(BaseModel):\n",
    "    questions: list[ClarifyingQuestion] = Field(description=\"List of clarifying questions\")\n",
    "\n",
    "clarifying_agent = Agent  (\n",
    "    name=\"ClarifyingAgent\",\n",
    "    instructions=CLARIFYING_INSTRUCTIONS,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=ClarifyingQuestions,  \n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03b24203",
   "metadata": {},
   "source": [
    "### Create a Search Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3eb00f6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "SEARCH_INSTRUCTIONS = \"You are a research assistant. Given a search term, you search the web for that term and \\\n",
    "produce a concise summary of the results. The summary must be 2-3 paragraphs and less than 300 \\\n",
    "words. Capture the main points. Write succinctly, no need to have complete sentences or good \\\n",
    "grammar. This will be consumed by someone synthesizing a report, so it's vital you capture the \\\n",
    "essence and ignore any fluff. Do not include any additional commentary other than the summary itself.\"\n",
    "\n",
    "search_agent = Agent(\n",
    "    name=\"Search agent\",\n",
    "    instructions=SEARCH_INSTRUCTIONS,\n",
    "    tools=[WebSearchTool(search_context_size=\"low\")],\n",
    "    model=\"gpt-4o-mini\",\n",
    "    model_settings=ModelSettings(tool_choice=\"required\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29b58222",
   "metadata": {},
   "source": [
    "### Create a Planner Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2242f4e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "HOW_MANY_SEARCHES = 3\n",
    "\n",
    "PLANNER_INSTRUCTIONS = f\"\"\"You are a helpful research assistant. Given a query and optional clarifications, \n",
    "come up with a set of web searches to perform to best answer the query. Output {HOW_MANY_SEARCHES} terms to query for.\n",
    "If clarifications are provided, use them to make your searches more targeted and relevant.\"\"\"\n",
    "\n",
    "class WebSearchItem(BaseModel):\n",
    "    reason: str = Field(description=\"Your reasoning for why this search is important to the query.\")\n",
    "    query: str = Field(description=\"The search term to use for the web search.\")\n",
    "\n",
    "class WebSearchPlan(BaseModel):\n",
    "    searches: list[WebSearchItem] = Field(description=\"A list of web searches to perform to best answer the query.\")\n",
    "\n",
    "planner_agent = Agent(\n",
    "    name=\"PlannerAgent\",\n",
    "    instructions=PLANNER_INSTRUCTIONS,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=WebSearchPlan,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c46f8d33",
   "metadata": {},
   "source": [
    "### Create a Notification Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "624fa7c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "@function_tool\n",
    "def send_research_notification(report_content: str) -> Dict[str, str]:\n",
    "    \"\"\" Send a research report as a Pushover notification \"\"\"\n",
    "    url = \"https://api.pushover.net/1/messages.json\"\n",
    "    \n",
    "    import re\n",
    "    clean_content = re.sub(r'#{1,6}\\s*', '', report_content)\n",
    "    clean_content = re.sub(r'\\*\\*(.*?)\\*\\*', r'\\1', clean_content)\n",
    "    clean_content = re.sub(r'\\*(.*?)\\*', r'\\1', clean_content)\n",
    "    clean_content = re.sub(r'\\n\\n+', '\\n\\n', clean_content)\n",
    "    \n",
    "    if len(clean_content) > 900:\n",
    "        clean_content = clean_content[:900] + \"...\\n\\n[Report truncated - see full version in trace]\"\n",
    "    \n",
    "    data = {\n",
    "        \"token\": os.environ.get('PUSHOVER_TOKEN'),\n",
    "        \"user\": os.environ.get('PUSHOVER_USER'),\n",
    "        \"message\": clean_content,\n",
    "        \"title\": \"Research Report Complete\"\n",
    "    }\n",
    "    \n",
    "    response = requests.post(url, data=data)\n",
    "    if response.status_code == 200:\n",
    "        return {\"status\": \"success\", \"message\": \"Research report notification sent successfully\"}\n",
    "    else:\n",
    "        return {\"status\": \"error\", \"message\": f\"Error {response.status_code}: {response.text}\"}\n",
    "\n",
    "notification_agent = Agent(\n",
    "    name=\"Notification agent\",\n",
    "    instructions=\"\"\"You are able to send a research report notification via Pushover.\n",
    "    You will be provided with a detailed markdown report. You should use your tool to send one notification \n",
    "    with the report content. The notification should be clean and readable.\"\"\",\n",
    "    tools=[send_research_notification],\n",
    "    model=\"gpt-4o-mini\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6184c9c4",
   "metadata": {},
   "source": [
    "### Create A Writer Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cf761d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "WRITER_INSTRUCTIONS = (\n",
    "    \"You are a senior researcher tasked with writing a cohesive report for a research query. \"\n",
    "    \"You will be provided with the original query, any clarifications provided by the user, and some initial research done by a research assistant.\\n\"\n",
    "    \"You should first come up with an outline for the report that describes the structure and \"\n",
    "    \"flow of the report. Then, generate the report and return that as your final output.\\n\"\n",
    "    \"The final output should be in markdown format, and it should be lengthy and detailed. Aim \"\n",
    "    \"for 5-10 pages of content, at least 1000 words.\"\n",
    ")\n",
    "\n",
    "class ReportData(BaseModel):\n",
    "    short_summary: str = Field(description=\"A short 2-3 sentence summary of the findings.\")\n",
    "    markdown_report: str = Field(description=\"The final report\")\n",
    "    follow_up_questions: list[str] = Field(description=\"Suggested topics to research further\")\n",
    "\n",
    "writer_agent = Agent(\n",
    "    name=\"WriterAgent\",\n",
    "    instructions=WRITER_INSTRUCTIONS,\n",
    "    model=\"gpt-4o-mini\",\n",
    "    output_type=ReportData,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0910bb12",
   "metadata": {},
   "source": [
    "### Core Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e4e6d5b",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def get_clarifying_questions(query: str):\n",
    "    \"\"\" Use the clarifying agent to generate questions about the query \"\"\"\n",
    "    print(\"Generating clarifying questions...\")\n",
    "    result = await Runner.run(clarifying_agent, f\"Research query: {query}\")\n",
    "    print(f\"Generated {len(result.final_output.questions)} questions\")\n",
    "    return result.final_output\n",
    "\n",
    "async def plan_searches(query: str, clarifications: str = \"\"):\n",
    "    \"\"\" Use the planner_agent to plan which searches to run for the query \"\"\"\n",
    "    print(\"Planning searches...\")\n",
    "    input_text = f\"Query: {query}\"\n",
    "    if clarifications:\n",
    "        input_text += f\"\\n\\nClarifications from user:\\n{clarifications}\"\n",
    "    result = await Runner.run(planner_agent, input_text)\n",
    "    print(f\"Will perform {len(result.final_output.searches)} searches\")\n",
    "    return result.final_output\n",
    "\n",
    "async def perform_searches(search_plan: WebSearchPlan):\n",
    "    \"\"\" Call search() for each item in the search plan \"\"\"\n",
    "    print(\"Searching...\")\n",
    "    tasks = [asyncio.create_task(search(item)) for item in search_plan.searches]\n",
    "    results = await asyncio.gather(*tasks)\n",
    "    print(\"Finished searching\")\n",
    "    return results\n",
    "\n",
    "async def search(item: WebSearchItem):\n",
    "    \"\"\" Use the search agent to run a web search for each item in the search plan \"\"\"\n",
    "    input = f\"Search term: {item.query}\\nReason for searching: {item.reason}\"\n",
    "    result = await Runner.run(search_agent, input)\n",
    "    return result.final_output\n",
    "\n",
    "async def write_report(query: str, clarifications: str, search_results: list[str]):\n",
    "    \"\"\" Use the writer agent to write a report based on the search results\"\"\"\n",
    "    print(\"Thinking about report...\")\n",
    "    input = f\"Original query: {query}\\n\\nClarifications: {clarifications}\\n\\nSummarized search results: {search_results}\"\n",
    "    result = await Runner.run(writer_agent, input)\n",
    "    print(\"Finished writing report\")\n",
    "    return result.final_output\n",
    "\n",
    "async def send_notification(report: ReportData):\n",
    "    \"\"\" Use the notification agent to send a notification with the report \"\"\"\n",
    "    print(\"Sending notification...\")\n",
    "    result = await Runner.run(notification_agent, report.markdown_report)\n",
    "    print(\"Notification sent\")\n",
    "    return report\n",
    "\n",
    "### Formatting Functions\n",
    "def format_questions(questions: ClarifyingQuestions) -> str:\n",
    "    \"\"\" Format clarifying questions for display \"\"\"\n",
    "    lines = [\n",
    "        \"Great! I'd like to ask a few clarifying questions to make the research more targeted:\\n\"\n",
    "    ]\n",
    "    \n",
    "    for i, q in enumerate(questions.questions, 1):\n",
    "        lines.append(f\"**{i}. {q.question}**\")\n",
    "    \n",
    "    lines.append(\"\\nPlease answer these questions (you can answer all at once or just provide key points).\")\n",
    "    return \"\\n\".join(lines)\n",
    "\n",
    "### Check if user is answering clarifying questions\n",
    "def is_answering_questions(history: List[Tuple[str, str]]) -> bool:\n",
    "    \"\"\" Check if user is answering clarifying questions \"\"\"\n",
    "    return bool(history and history[-1][1] and \"clarifying questions\" in history[-1][1].lower())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "159d40b1",
   "metadata": {},
   "source": [
    "### Main Research Flow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10cf0a52",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def run_research(query: str, clarifications: str):\n",
    "    \"\"\" Main research flow with clarifications \"\"\"\n",
    "    with trace(\"Research trace\"):\n",
    "        print(\"Starting research...\")\n",
    "        search_plan = await plan_searches(query, clarifications)\n",
    "        search_results = await perform_searches(search_plan)\n",
    "        report = await write_report(query, clarifications, search_results)\n",
    "        await send_notification(report)  \n",
    "        print(\"Hooray!\")\n",
    "        return report.markdown_report"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd8c1db1",
   "metadata": {},
   "source": [
    "### Gradio UI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b0b8a1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "async def chat(message: str, history: List[Tuple[str, str]]):\n",
    "    \"\"\" Handle chat messages through the research flow \"\"\"\n",
    "    if is_answering_questions(history):\n",
    "        # Run research with user's answers\n",
    "        original_query = history[0][0]\n",
    "        yield \"Perfect! Starting the research now... This may take a minute.\\n\\n⏳ Planning searches...\"\n",
    "        report = await run_research(original_query, message)\n",
    "        yield f\"✅ Research complete!\\n\\n{report}\\n\\n---\\n\\nFeel free to ask me to research another topic!\"\n",
    "    else:\n",
    "        # Ask clarifying questions for new topic\n",
    "        questions = await get_clarifying_questions(message)\n",
    "        yield format_questions(questions)\n",
    "\n",
    "\n",
    "ui = gr.ChatInterface(\n",
    "    fn=chat,\n",
    "    title=\"🔍 Deep Research Assistant\",\n",
    "    description=\"Tell me what you'd like to research, and I'll ask clarifying questions before diving deep!\",\n",
    "    examples=[\n",
    "        \"Latest AI Agent frameworks in 2025\",\n",
    "        \"Best practices for RAG systems\",\n",
    "        \"State of quantum computing\"\n",
    "    ],\n",
    "    theme=gr.themes.Soft(primary_hue=\"blue\"),\n",
    ")\n",
    "\n",
    "ui.launch(inbrowser=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
